package in.ac.iitb.cse.nlp.wordprediction;

import in.ac.iitb.cse.nlp.postagger.data.DataMaps;
import in.ac.iitb.cse.nlp.postagger.data.Matrix;
import in.ac.iitb.cse.nlp.postagger.data.Word;
import in.ac.iitb.cse.nlp.poswordpredictor.Constants;

import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

public class WordMatrix implements Matrix {

	// Internal map which stores the word matrix
	private Map<Word, HashMap<Word, Integer>> wordTable = new HashMap<Word, HashMap<Word, Integer>>();

	public WordMatrix() {

	}

	@Override
	public void update(String previous, String current) {
		Word previousWord = DataMaps.words.get(previous);
		if (previousWord == null) {
			previousWord = new Word(previous);
			DataMaps.words.put(previous, previousWord);
		} else {
			previousWord.incrementCount();
		}
		HashMap<Word, Integer> map = wordTable.get(previousWord);
		if (map == null) {
			map = new HashMap<Word, Integer>();
			wordTable.put(previousWord, map);
		}
		Word currentWord = DataMaps.words.get(current);
		if (currentWord == null) {
			currentWord = new Word(current);
			DataMaps.words.put(current, currentWord);
		}
		Integer count = map.get(currentWord);
		if (count == null) {
			map.put(currentWord, 1);
		} else {
			map.put(currentWord, count + 1);
		}

	}

	@Override
	public int getCount(String previous, String current) {
		Word previousWord = DataMaps.words.get(previous);
		if (previousWord != null) {
			HashMap<Word, Integer> hashMap = wordTable.get(previousWord);
			if (hashMap != null) {
				Word currentWord = DataMaps.words.get(current);
				if (currentWord != null) {
					Integer integer = hashMap.get(currentWord);
					if (integer != null) {
						return integer;
					}
				}
			}
		}
		return 0;
	}

	@Override
	public double getProbability(String previous, String current) {
		int countOfCurrentGivenPrev = getCount(previous, current);
		if (countOfCurrentGivenPrev != 0) {
			Word previousWord = DataMaps.words.get(previous);
			return ((double) countOfCurrentGivenPrev / (double) previousWord
					.getCount());
		}
		return 0;
	}

	List<String> getSuitablePrediction(String wordString) {
		Word word = DataMaps.words.get(wordString);
		List<String> predictions = new ArrayList<String>();
		if (word != null) {
			HashMap<Word, Integer> hashMap = wordTable.get(word);
			List<Prediction> temp = new ArrayList<Prediction>();
			if (hashMap != null) {
				for (Word key : hashMap.keySet()) {
					temp.add(new Prediction(key.getString(), getProbability(
							word.getString(), key.getString())));
				}
			}
			Collections.sort(temp);
			for (int i = 0; i < temp.size()
					&& i < Constants.MAX_NUM_OF_PREDICTONS; i++) {
				predictions.add(temp.get(i).word);
			}
		}
		return predictions;
	}

	@Override
	public String toString() {
		StringBuffer buffer = new StringBuffer();
		for (Word prev : wordTable.keySet()) {
			for (Word current : wordTable.get(prev).keySet()) {
				buffer.append("(" + prev + "," + current + ")="
						+ prev.getCount() + "," + current.getCount() + ","
						+ wordTable.get(prev).get(current) + ":"
						+ getProbability(prev.getString(), current.getString()));
			}
			buffer.append("\n");
		}
		return buffer.toString();
	}

}

class Prediction implements Comparable<Prediction> {

	String word;
	Double probability;

	Prediction(String word, Double probability) {
		this.word = word;
		this.probability = probability;
	}

	@Override
	public int compareTo(Prediction prediction) {
		// We want to sort in descending order
		if (probability < prediction.probability) {
			return 1;
		} else if (probability > prediction.probability) {
			return -1;
		} else {
			return 0;
		}
	}

}
